import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Circle

class Visualization:
    """
    可视化功能
    """
    def __init__(self):
        self.fig = None
        self.axes = None
        
    def plot_trajectory(self, results, errors=None, save_path=None):
        """
        绘制轨迹图
        results: 仿真结果
        errors: 误差计算结果
        save_path: 保存路径
        """
        # 提取数据
        time_points = results['time_points']
        states_history = results['states_history']
        reference_trajectory = results['reference_trajectory']
        
        # 创建图形
        self.fig, self.axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # 绘制轨迹 - 增加线宽和对比度
        self.axes[0, 0].plot(reference_trajectory[:, 0], reference_trajectory[:, 1], 
                              'b-', linewidth=2.5, label='参考轨迹')
        self.axes[0, 0].plot(states_history[:, 0], states_history[:, 1], 
                              'r--', linewidth=2, label='实际轨迹')
        
        # 标记起点和终点
        self.axes[0, 0].scatter(reference_trajectory[0, 0], reference_trajectory[0, 1], 
                                 c='g', s=100, marker='o', label='起点')
        self.axes[0, 0].scatter(reference_trajectory[-1, 0], reference_trajectory[-1, 1], 
                                 c='m', s=100, marker='*', label='终点')
        
        self.axes[0, 0].set_xlabel('X位置', fontsize=12)
        self.axes[0, 0].set_ylabel('Y位置', fontsize=12)
        self.axes[0, 0].set_title('轨迹跟踪', fontsize=14)
        self.axes[0, 0].legend(fontsize=12)
        self.axes[0, 0].grid(True)
        self.axes[0, 0].axis('equal')
        
        # 添加当前时间和RMSE信息
        if errors is not None:
            rmse_text = f"RMSE: {errors['rmse']:.4f}"
            max_error_text = f"最大误差: {errors['max_error']:.4f}"
            time_text = f"时间: {time_points[-1]:.1f}s"
            
            self.axes[0, 0].text(0.02, 0.95, time_text, transform=self.axes[0, 0].transAxes, 
                                 fontsize=12, verticalalignment='top')
            self.axes[0, 0].text(0.02, 0.90, rmse_text, transform=self.axes[0, 0].transAxes, 
                                 fontsize=12, verticalalignment='top')
            self.axes[0, 0].text(0.02, 0.85, max_error_text, transform=self.axes[0, 0].transAxes, 
                                 fontsize=12, verticalalignment='top')
        
        # 绘制位置随时间变化 - 增加线宽和对比度
        self.axes[0, 1].plot(time_points[:len(states_history)], reference_trajectory[:len(states_history), 0], 
                             'b-', linewidth=2, label='X参考')
        self.axes[0, 1].plot(time_points[:len(states_history)], states_history[:, 0], 
                             'r--', linewidth=1.5, label='X实际')
        self.axes[0, 1].plot(time_points[:len(states_history)], reference_trajectory[:len(states_history), 1], 
                             'g-', linewidth=2, label='Y参考')
        self.axes[0, 1].plot(time_points[:len(states_history)], states_history[:, 1], 
                             'm--', linewidth=1.5, label='Y实际')
        self.axes[0, 1].set_xlabel('时间(s)', fontsize=12)
        self.axes[0, 1].set_ylabel('位置', fontsize=12)
        self.axes[0, 1].set_title('位置随时间变化', fontsize=14)
        self.axes[0, 1].legend(fontsize=10)
        self.axes[0, 1].grid(True)
        
        # 绘制速度随时间变化 - 增加线宽和对比度
        self.axes[1, 0].plot(time_points[:len(states_history)], reference_trajectory[:len(states_history), 2], 
                             'b-', linewidth=2, label='Vx参考')
        self.axes[1, 0].plot(time_points[:len(states_history)], states_history[:, 2], 
                             'r--', linewidth=1.5, label='Vx实际')
        self.axes[1, 0].plot(time_points[:len(states_history)], reference_trajectory[:len(states_history), 3], 
                             'g-', linewidth=2, label='Vy参考')
        self.axes[1, 0].plot(time_points[:len(states_history)], states_history[:, 3], 
                             'm--', linewidth=1.5, label='Vy实际')
        self.axes[1, 0].set_xlabel('时间(s)', fontsize=12)
        self.axes[1, 0].set_ylabel('速度', fontsize=12)
        self.axes[1, 0].set_title('速度随时间变化', fontsize=14)
        self.axes[1, 0].legend(fontsize=10)
        self.axes[1, 0].grid(True)
        
        # 如果有误差数据，绘制误差
        if errors is not None:
            self.axes[1, 1].plot(time_points[:len(errors['distance_errors'])], errors['distance_errors'], 
                                 'k-', linewidth=2)
            self.axes[1, 1].set_xlabel('时间(s)', fontsize=12)
            self.axes[1, 1].set_ylabel('跟踪误差', fontsize=12)
            self.axes[1, 1].set_title(f'跟踪误差 (RMSE: {errors["rmse"]:.4f}, 最大误差: {errors["max_error"]:.4f})', 
                                      fontsize=14)
            self.axes[1, 1].grid(True)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图像
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"轨迹图已保存到: {save_path}")
            
        plt.show()
        
    def create_animation(self, results, save_path=None):
        """
        创建动画
        results: 仿真结果
        save_path: 保存路径
        """
        # 提取数据
        states_history = results['states_history']
        reference_trajectory = results['reference_trajectory']
        time_points = results['time_points'][:len(states_history)]
        
        # 创建图形
        fig, ax = plt.subplots(figsize=(10, 8))
        
        # 设置坐标轴
        x_min = min(np.min(states_history[:, 0]), np.min(reference_trajectory[:, 0])) - 2
        x_max = max(np.max(states_history[:, 0]), np.max(reference_trajectory[:, 0])) + 2
        y_min = min(np.min(states_history[:, 1]), np.min(reference_trajectory[:, 1])) - 2
        y_max = max(np.max(states_history[:, 1]), np.max(reference_trajectory[:, 1])) + 2
        
        ax.set_xlim(x_min, x_max)
        ax.set_ylim(y_min, y_max)
        ax.set_xlabel('X位置', fontsize=12)
        ax.set_ylabel('Y位置', fontsize=12)
        ax.set_title('双积分小车LQR轨迹跟踪控制', fontsize=14)
        ax.grid(True)
        
        # 绘制参考轨迹 - 使用更明显的线条
        ax.plot(reference_trajectory[:, 0], reference_trajectory[:, 1], 
                'b-', alpha=0.7, linewidth=2.5, label='参考轨迹')
        
        # 初始化实际轨迹线
        actual_line, = ax.plot([], [], 'r-', linewidth=2, label='实际轨迹')
        
        # 小车表示 - 增加尺寸和颜色对比度
        vehicle = Circle((0, 0), radius=0.3, fill=True, color='red', alpha=0.8)
        ax.add_patch(vehicle)
        
        # 参考点表示 - 增加尺寸和颜色对比度
        reference_point = Circle((0, 0), radius=0.3, fill=True, color='blue', alpha=0.8)
        ax.add_patch(reference_point)
        
        # 文本 - 增加字体大小和背景
        time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes, fontsize=12,
                           bbox=dict(facecolor='white', alpha=0.7))
        error_text = ax.text(0.02, 0.90, '', transform=ax.transAxes, fontsize=12,
                            bbox=dict(facecolor='white', alpha=0.7))
        
        ax.legend(fontsize=12)
        
        # 初始化函数
        def init():
            actual_line.set_data([], [])
            vehicle.center = (states_history[0, 0], states_history[0, 1])
            reference_point.center = (reference_trajectory[0, 0], reference_trajectory[0, 1])
            time_text.set_text('')
            error_text.set_text('')
            return actual_line, vehicle, reference_point, time_text, error_text
        
        # 更新函数
        def update(frame):
            # 更新实际轨迹
            actual_line.set_data(states_history[:frame+1, 0], states_history[:frame+1, 1])
            
            # 更新小车位置
            vehicle.center = (states_history[frame, 0], states_history[frame, 1])
            
            # 更新参考点位置
            ref_idx = min(frame, len(reference_trajectory)-1)
            reference_point.center = (reference_trajectory[ref_idx, 0], reference_trajectory[ref_idx, 1])
            
            # 计算当前误差
            error = np.sqrt((states_history[frame, 0] - reference_trajectory[ref_idx, 0])**2 + 
                           (states_history[frame, 1] - reference_trajectory[ref_idx, 1])**2)
            
            # 更新文本
            time_text.set_text(f'时间: {time_points[frame]:.1f}s')
            error_text.set_text(f'跟踪误差: {error:.4f}')
            
            return actual_line, vehicle, reference_point, time_text, error_text
        
        # 创建动画
        anim = FuncAnimation(fig, update, frames=len(states_history),
                              init_func=init, blit=True, interval=50)
        
        # 保存动画
        if save_path:
            anim.save(save_path, writer='pillow', fps=20)
            print(f"动画已保存到: {save_path}")
            
        plt.tight_layout()
        plt.show()
        
        return anim 